'''
Description:  情感识别 数据格式化 标签对齐
version: 
Author: 
Date: 2023-12-18 15:36:14
LastEditors: Please set LastEditors
LastEditTime: 2023-12-19 14:28:40
'''

#!/usr/bin/python
# -*- coding: UTF-8 -*-
 
from xml.dom.minidom import parse
import xml.dom.minidom
import json

file_path = "data"
none_label = 'neutral'
labels = [none_label,'sadness','like','anger','disgust','fear','happiness','surprise']

#计算data中各个label的样本数
def countLabels(data,labels):
  print("语料数量:",len(data))
  for label in labels:
    newlist = list(d for d in data if d['label'] == labels.index(label))
    print(label+":",len(newlist))

# 处理微博4000条8类情感数据
def propcess1():
  # 使用minidom解析器打开 XML 文档
  DOMTree = xml.dom.minidom.parse(file_path + "/weibo-emotion-4000.xml")
  collection = DOMTree.documentElement
  dialogs = collection.getElementsByTagName("weibo")

  data = []

  for dialog in dialogs:
    if dialog.getElementsByTagName("sentence"):
        sentences = dialog.getElementsByTagName("sentence")
        for s in sentences:
          if(len(s.childNodes)):
            text = s.childNodes[0].data.strip()
            emotion = []
            if s.getAttribute("emotion_tag")  == 'Y':
              emotion.append(s.getAttribute("emotion-1-type"))
              if s.getAttribute("emotion-2-type") != 'none':
                emotion.append(s.getAttribute("emotion-2-type"))
            if s.getAttribute("emotion_tag")  == 'N':
              emotion.append(none_label)
            data.append({
              'text':text.strip(),
              'label':labels.index(emotion[0]),
              'emotion':emotion
            })

  countLabels(data=data,labels=labels)

  with open(file_path + "/preprocess/weibo-emotion-4000.json","w", encoding='utf-8') as f:
    json.dump(data,f, ensure_ascii=False)
    print("\n加载入文件完成...")

#emotion-class-6.json
def propcess2():
  # 其他（none), 喜好(Like)，悲伤(Sad)，厌恶(Disgust)，愤怒(Anger)，高兴（Happiness）六类，依次标号为0到5
  _labels = [none_label,'like','sadness','disgust','anger','happiness']
  
  with open(file_path + '/emotion-class-6.json', encoding="utf-8") as f:
    raw_data = json.load(f)
  data = []
  for d in raw_data:
    data.append({
      'text':d[0].replace(" ","").strip(),
      'label':labels.index(_labels[d[1]]),
      'emotion':[_labels[d[1]]]
    })

  countLabels(data=data,labels=labels)

  with open(file_path + "/preprocess/emotion-class-6-"+str(len(data))+".json","w", encoding='utf-8') as f:
    json.dump(data,f, ensure_ascii=False)
    print("\n加载入文件完成...")

# 3-SMP2020微博情绪分类技术评测 usual_train
def propcess3():
  _labels = [none_label,'sad','','angry','','fear','happy','surprise']

  with open(file_path + '/usual_train.json', encoding="utf-8") as f:
    raw_data = json.load(f)
  data = []
  for d in raw_data:
    emotion = [labels[_labels.index(d['label'])]]
    data.append({
      'text':d['content'].strip(),
      'label':labels.index(emotion[0]),
      'emotion':emotion
    })

  countLabels(data=data,labels=labels)


  with open(file_path + "/preprocess/usual-train-class-6-"+str(len(data))+".json","w", encoding='utf-8') as f:
    json.dump(data,f, ensure_ascii=False)
    print("\n加载入文件完成...")


propcess1()
propcess2()
propcess3()
